/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution
{
public:
    vector<int> distanceK(TreeNode *root, TreeNode *target, int k)
    {
        if (k == 0)
        {
            return {target->val};
        }
        unordered_map<int, vector<int>> adjTable;
        setupAdjTable(root, adjTable);
        unordered_set<int> visited;
        queue<int> q;
        visited.insert(target->val);
        q.push(target->val);
        int dist = 0;
        while (!q.empty())
        {
            int qSize = q.size();
            for (int i = 0; i < qSize; ++i)
            {
                int node = q.front();
                q.pop();
                for (int adjNode : adjTable[node])
                {
                    if (!visited.count(adjNode))
                    {
                        q.push(adjNode);
                        visited.insert(adjNode);
                    }
                }
            }
            ++dist;
            if (dist == k)
            {
                vector<int> result;
                while (!q.empty())
                {
                    result.push_back(q.front());
                    q.pop();
                }
                return result;
            }
        }
        return {};
    }

    void setupAdjTable(TreeNode *tree, unordered_map<int, vector<int>> &adjTable)
    {
        if (tree->left != nullptr)
        {
            adjTable[tree->val].push_back(tree->left->val);
            adjTable[tree->left->val].push_back(tree->val);
            setupAdjTable(tree->left, adjTable);
        }
        if (tree->right != nullptr)
        {
            adjTable[tree->val].push_back(tree->right->val);
            adjTable[tree->right->val].push_back(tree->val);
            setupAdjTable(tree->right, adjTable);
        }
    }
};